#!/usr/bin/env python


import atexit
import socket
import struct
import pickle
from Queue import Queue
from threading import Thread, currentThread
from random import randint


ALLOW_PEERS = 10
INSTANCES = set()


def stop_all_instances():
    """
    Stops all running instances
    of RemoteQueue. 
    """
    for i in INSTANCES:
        i._stop()

atexit.register(stop_all_instances)


class RemoteQueue(Queue):
    """
        An implementation of the Queue class
    for safely sharing objects not just between
    threads but also between processes and event-
    ually between machines.
    
        Every time an instance is created a refer-
    ence is placed in the global INSTANCES set in
    order for safe exit at interpreter shutdown.
    Then the instance searches over a port range
    for other active instances and places them in
    its set of peers. It then broadcasts its
    presence to all known peers and starts its
    listener thread.
    
        Every time an object is put on the queue
    the instance generates a random ID and broadcasts
    the object along with the ID to a specific sub-
    section of its peers. The random ID ensures that
    when the instance receives its message back from
    a peer it knows to ignore it.
    
        The listener thread handles puts from other
    RemoteQueue instances and also adding/removing
    instances from its local set of peers. Every
    incoming message is handled by an accepter
    thread so that multiple puts can be handled
    simultaneously.
    """
    def __init__(self, debug=False, port=24000):
        """
        debug = boolean controlling the printing of
            important information
        port = the port number of the 'server' or
            the first peer in the group
        """
        INSTANCES.add(self)
        self._accepting = False
      	self.addr = socket.gethostbyname(socket.gethostname())
        self._debug = debug
        self._bound = False
        self._peers = set()
        self._msg_ids = set()
        self._find_peers(self.addr, port)
        self._bcast_peers('!!!', 
                          self._peers | set([(self.addr,self.port)]), 
                          self._gen_rand_id())

        self.rsocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.rsocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.rsocket.bind((self.addr,self.port))
        self.rsocket.listen(5)

        self._acceptors = set()
        self._thread = Thread(target=self._listen)
        self._thread.setDaemon(True)
        self._thread.start()

	Queue.__init__(self)

    def _gen_rand_id(self):
        """
        Generates a random ID for use in
        tracking messages.
        """
        return randint(2**4,2**30)

    def _find_peers(self, addr, port):
        """
        Port scans within a range of ports
        for peers with active connections.
        The first inactive connection it finds
        becomes the instances port.
        
        Note: an active connection does not
        necessarily mean a RemoteQueue
        instance is on the other side.
        
        addr = for now this is just the local
            loopback
        port = port to start search from
        """
        for p in range(port,port+ALLOW_PEERS):
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            err = s.connect_ex((addr, p))
            if not bool(err):
                s.send(struct.pack(">L", 0))
                self._peers.add((addr,p))
            elif not self._bound:
                self.port = p
                self._bound = True
            s.close()

    def _bcast_peers(self, cmd, obj, id):
        """
        Broadcasts given message to all known
        peers.

        cmd = { '!!!' : new peer available,
                '---' : peer has exited,
                'put' : put given object
                        in Queue }
        obj = any hashable Python object
        id = unique identifier, preferably gen-
             erated by _gen_rand_id
        """
        for addr in self._peers.copy():
            msg = cmd+':'+pickle.dumps(obj)+':'+str(id)
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            try:
                s.connect(addr)
                s.send(struct.pack(">L", len(msg)))
                s.send(msg)
            except socket.error:
                if self._debug: print "Connect failed, trying again..."
                #self._bcast_peers(cmd, obj, id)
                return
            s.close()

    def _accept(self, client, address):
        """
        Handles incoming messages, usually
        called by the listener thread upon
        receiving a client and called in a
        dedicated thread.
        """
        header = ""
        while len(header) <> 4:
            header += client.recv(1)
        msg_len = struct.unpack(">L", header)[0]
        msg = client.recv(msg_len).split(':')
        if len(msg) > 1:
            id = int(msg[2])
            cmd = msg[0]
            obj = pickle.loads(msg[1])
            if id not in self._msg_ids:
                self._msg_ids.add(id)
                if cmd == '!!!':
                    self._peers |= obj
                    self._peers -= set([(self.addr,self.port)])
                elif cmd == '---':
                    self._peers -= obj
                elif cmd == 'put':
                    self.put(obj,id)
        client.close()
        #self._acceptors.remove(currentThread())
                
    def _listen(self):
        """
        Started at initiallization and stopped
        by the _stop method, it waits for mess-
        ages from clients and starts accepter
        threads to handle them.
        """
        self._listening = True
        while self._listening:
            self._accepting = True
            client, address = self.rsocket.accept()
            acceptor = Thread(target=self._accept,
                              args=(client,address))
            self._acceptors.add(acceptor)
            acceptor.setDaemon(False)
            acceptor.start()

    def _connect_to_self(self):
        """
        Used by the _stop method to unblock
        the blocking socket.accept() call in
        order to safely kill the listener
        thread.
        """
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((self.addr, self.port))
        s.send(struct.pack(">L", 0))
        s.send("0")
        s.close()

    def _bcast_exit(self):
        """
        Broadcasts exit status to all known peers,
        it is the first method called by _stop.
        """
        self._bcast_peers('---', 
                          set([(self.addr,self.port)]),
                          self._gen_rand_id())
        
    def _stop(self):
        """
        Safely kills the listener thread, called
        on all active instances by stop_all_instances
        function at interpreter shutdown.
        """
        while not self._accepting:
            pass
        self._bcast_exit()
        self._listening = False
        self._connect_to_self()
        self._thread.join()
        for a in self._acceptors:
            a.join()

    def add_peer(self, peer):
        """
        Given a socket client tuple (addr, port)
        add it to the set of known peers and broadcast
        the set to all known peers.
        """
	self._find_peers(peer[0], peer[1])
        self._bcast_peers('!!!', self._peers | \
                          set([(self.addr,self.port)]), self.port)

    def put(self, obj, id=None):
        """
        Overriden Queue.put that appends a hopefully
        unique ID to all puts and broadcasts it to
        all known peers.
        """
        Queue.put(self, obj)
        if not id:
            id = self._gen_rand_id()
            self._msg_ids.add(id)
        self._bcast_peers('put', obj, id)


remoteQueue = RemoteQueue # for backwards compatibility
